const Router = require('koa-router');
const { Ollama } = require('@langchain/ollama');
const { PromptTemplate } = require('@langchain/core/prompts');
const { StringOutputParser } = require('@langchain/core/output_parsers');
const { schema } = require('../lib/schema');
const { generatePrompt } = require('../prompts/generate');
const { summarizePrompt } = require('../prompts/summarize');
const { connectDB } = require('../lib/db');

const router = new Router();

// 初始化 Ollama 模型
const llm = new Ollama({
  model: 'qwen3:0.6b',
  temperature: 0,
  maxRetries: 2,
});

// 执行 MongoDB 查询
async function executeQuery(pipeline, collectionName = 'jobtickets') {
  const db = await connectDB();
  const collection = db.collection(collectionName);
  const cursor = collection.aggregate(pipeline);
  return await cursor.toArray();
}

// 校验聚合管道（防注入、非法字段）
function validatePipeline(pipeline) {
  const errors = [];
  pipeline.forEach((stage, i) => {
    if (stage.$lookup) {
      const from = stage.$lookup.from;
      if (!['types', 'departments'].includes(from)) {
        errors.push(`非法集合: ${from}`);
      }
    }
  });
  return { valid: errors.length === 0, errors };
}

// 主接口
router.post('/ask', async (ctx) => {
  const { question } = ctx.request.body;

  if (!question) {
    ctx.status = 400;
    ctx.body = { error: '缺少问题' };
    return;
  }

  try {
    // Step 1: 生成聚合查询
    const generateChain = PromptTemplate
      .fromTemplate(generatePrompt)
      .pipe(llm)
      .pipe(new StringOutputParser());

    const rawQuery = await generateChain.invoke({ schema, question });
    let pipeline;

    try {
      pipeline = JSON.parse(rawQuery.trim());
    } catch (e) {
      throw new Error('AI 返回非 JSON 查询: ' + rawQuery);
    }

    // 校验查询
    const { valid, errors } = validatePipeline(pipeline);
    if (!valid) {
      ctx.status = 400;
      ctx.body = { error: '查询校验失败', details: errors };
      return;
    }

    // Step 2: 执行查询
    const result = await executeQuery(pipeline);

    // Step 3: AI 播报结果
    const summarizeChain = PromptTemplate
      .fromTemplate(summarizePrompt)
      .pipe(llm)
      .pipe(new StringOutputParser());

    const answer = await summarizeChain.invoke({
      query: JSON.stringify(pipeline, null, 2),
      result: JSON.stringify(result, null, 2),
    });

    // 返回结果
    ctx.body = {
      question,
      query: pipeline,
      data: result,
      answer: answer.trim(),
    };
  } catch (err) {
    ctx.status = 500;
    ctx.body = { error: err.message };
  }
});

module.exports = router;
